#include <stdio.h>
#include <stdlib.h>
#include <getopt.h>
#include <limits.h>
#include <string.h>
#include <stdint.h>
#include <openssl/ec.h>
#include <openssl/evp.h>
#include <openssl/sm2.h>
#include <openssl/sm3.h>
#include <openssl/bn.h>

#include "crypto.h"
#include "session_data.h"
#include "key_cert.h"
#include "utils.h"
#include "ovmf_hash_table.h"
#include "cal_vm_digest.h"

typedef enum option_choice {
    OPT_ERR = -1,
    OPT_EOF = 0,
    OPT_HELP,
    OPT_VERBOSE,
    OPT_BIOS,
    OPT_BUILD_ID,
    OPT_VM_ID,
    OPT_KERNEL,
    OPT_CMDLINE,
    OPT_INITRD,
    OPT_VM_VERSION,
    OPT_VCPU_MODEL,
    OPT_VCPU_NUM,
    OPT_USER_ID,
    OPT_USER_DATA,
    OPT_GUEST_POLICY,
    OPT_GUEST_PRIV_KEY,
} option_choice_t;

char *vcpu_model[] = { "Dhyana", "Dhyana-v1", "Dhyana-v2", "Dhyana-v3", "Dharma", "Dharma-v1", "host"};

static struct option long_args_options[] = {
    {"help",          no_argument,       0, OPT_HELP},
    {"verbose",       no_argument,       0, OPT_VERBOSE},
    {"build",         required_argument, 0, OPT_BUILD_ID},
    {"bios",          required_argument, 0, OPT_BIOS},
    {"kernel",        required_argument, 0, OPT_KERNEL},
    {"cmdline",       required_argument, 0, OPT_CMDLINE},
    {"initrd",        required_argument, 0, OPT_INITRD},
    {"vm_id",         required_argument, 0, OPT_VM_ID},
    {"vm_version",    required_argument, 0, OPT_VM_VERSION},
    {"user_id",       required_argument, 0, OPT_USER_ID},
    {"user_data",     required_argument, 0, OPT_USER_DATA},
    {"cpu",           required_argument, 0, OPT_VCPU_MODEL},
    {"smp",           required_argument, 0, OPT_VCPU_NUM},
    {"guest_policy",  required_argument, 0, OPT_GUEST_POLICY},
    {"guest_private_key",  required_argument, 0, OPT_GUEST_PRIV_KEY},
    {0, 0, 0, 0}
};

const char* MASTER_LABEL = "csv-master-secret";
const char* KEK_LABEL    = "csv-kek";
const char* KIK_LABEL    = "csv-kik";

/**
 * @brief  Check if the vcpu model is in the range
 *
 * @param  [in]  model     The vcpu model string
 *
 * @returns 0 for success, -1 for failure
 */
int good_vcpu_model(char *model)
{
    int array_size = sizeof(vcpu_model) / sizeof(vcpu_model[0]);
    int n = 0;

    if (!model)
        return -1;

    for(n = 0; n < array_size; n++) {
        if (strcmp(model, vcpu_model[n]) == 0) {
            return 0;
        }
    }

    printf("Wrong cpu model, please choose vcpu model from :\n");
    for(n = 0; n < array_size; n++) {
        printf("%s ", vcpu_model[n]);
    }
    printf("\n");

    return -1;
}

/**
 * @brief  Usage
 *
 */

void usage(void)
{
    printf("Usage: program [options]\n");
    printf("Options:\n");
    printf("  --help,                         Display this summary\n");
    printf("  --verbose,                      Enable debug log print\n");
    printf("  --build <build_id>,             [optional] Input the build id of the Firmware\n");
    printf("  --bios <guest_bios_file>,       Input the virtual machine bios file\n");
    printf("  --guest_policy <policy>,        the virtual machine's policy\n");
    printf("  --guest_private_key <priv_key>, user's DH private key in PEM format\n");
    printf("  --kernel <guest_kernel_file>,   [optional] Input the virtual machine kernel file\n");
    printf("  --cmdline <guest_cmdline>,      [optional] Input the virtual machine command line file\n");
    printf("  --initrd <guest_initrd_file>,   [optional] Input the virtual machine initrd file\n");
    printf("  --vm_id <vm_id>,                [optional] Virtual machine id ,the length must be less than 16 bytes\n");
    printf("  --vm_version <version>,         [optional] Virtual machine version ,the length must be less than 16 bytes\n");
    printf("  --cpu <vcpu_model>,             [optional] Hygon vCPU model 'host' or (Dhyana{,-v1,-2,-v3},Dharma{,-v1}), required by CSV3\n");
    printf("  --smp <vcpu_num>,               [optional] The number of vCPUs of the guest, required by CSV3\n");
    printf("  --user_id <user_id>,            [optional] user's DH key id, length must be less than 254 bytes\n");
    printf("  --user_data <version>,          [optional] user's self-defined data, length must be less than 64 bytes\n");
}

/*
 *  Master Key:
 *   256-bit value, to derive KIK/KEK keys
 *
 *  Key Integrity Key(KIK):
 *   HMAC SHA-256, protect key's integrity
 *   256-bit derived from Master Key
 *
 *  Key Encryption Key(KEK):
 *   AES-128, protect key'ssecret
 *   128-bit derived from Master Key
 *
 *  Transport Integrity Key(TIK):
 *   HMAC SHA-256, protect transport integrity
 *   first half of 256-bit random value
 *
 *  Transport Encryption Key(TEK):
 *   AES-128, protect transport secret
 *   128-bit random value
 *
 *  This function will :
 *  1) generate a root key, which derives Master Key
 *
 *  2) generate the TIK/TEK randomly, the 2 keys will be used to
 *  protect the data transportation between user and hardware
 *
 *  3) Derive the KIK/KEK from Master Key, the 2 keys will wrap
 *     the TIK/TEK
 *
 *  4) fill the tek_tik related to session and calculate mac
 */

int  generate_tk(session_buf* out_session_data,
                 tek_tik   *tek_tik_value,
                 nonce_256 *ms,
                 uint32_t build_id)
{
    int           ret            = -1;
    nonce_256     root_key       = {0};
    nonce_256     master_key     = {0};
    aes_128_key   kek_key        = {0};
    hmac_sha_256  kik_key        = {0};
    unsigned char tek_tik_random[TEK_TIK_LEN]      = {0};

    uint8_t       hmac_buffer[sizeof(out_session_data->wrap_iv) + sizeof(out_session_data->wrap_tk)];

    gen_random_bytes(out_session_data->nonce, sizeof(nonce_128));
    gen_random_bytes(out_session_data->wrap_iv, sizeof(iv_128));
    gen_random_bytes(tek_tik_random, TEK_TIK_LEN);

    memcpy(tek_tik_value, tek_tik_random, sizeof(tek_tik_random));

    /* root key is used to derived maste key and so on */
    gen_random_bytes(root_key, sizeof(root_key));

    if (kdf_gen_key(root_key, sizeof(root_key),
                    (unsigned char*)MASTER_LABEL, (uint32_t)strlen(MASTER_LABEL),
                    out_session_data->nonce, sizeof(out_session_data->nonce),
                    master_key, sizeof(master_key)) != 0) {
        printf("fail to generate master key\n");
        goto finish;
    }

    if (kdf_gen_key(master_key, sizeof(master_key),
                   (unsigned char*)KIK_LABEL, (uint32_t)strlen(KIK_LABEL),
                   NULL, 0,
                   kik_key, sizeof(kik_key)) != 0) {
        printf("fail to generate kik key\n");
        goto finish;
    }

    if (kdf_gen_key(master_key, sizeof(master_key),
                    (unsigned char*)KEK_LABEL, (uint32_t)strlen(KEK_LABEL),
                    NULL, 0,
                    kek_key, sizeof(kek_key))) {
        printf("fail to generate kek key\n");
        goto finish;
    }

    /* encrypt the tek_tik by kek, the cipher is stored in wrap_tk*/
    if (sm4_ctr128_encrypt(kek_key, out_session_data->wrap_iv,
                           tek_tik_random, TEK_TIK_LEN,
                       (unsigned char*)&out_session_data->wrap_tk) != 0) {
        printf("fail to encrypt the tek_tik\n");
        goto finish;
    }


    /* create the hmac for wrap_tk by kik */
    if (build_id <= 1544) {
        sm3_hmac((unsigned char*)(out_session_data->wrap_tk.tek),
                 sizeof(out_session_data->wrap_tk),
                 kik_key, sizeof(kik_key),
                 out_session_data->wrap_mac);
    } else {
        memcpy(hmac_buffer, &out_session_data->wrap_iv, sizeof(out_session_data->wrap_iv));
        memcpy(hmac_buffer + sizeof(out_session_data->wrap_iv),
               &out_session_data->wrap_tk, sizeof(out_session_data->wrap_tk));
        sm3_hmac(hmac_buffer, sizeof(hmac_buffer), kik_key, sizeof(kik_key),
                 out_session_data->wrap_mac);
    }

    memcpy(ms, master_key, sizeof(master_key));

    /* temporarily save the master key in session buffer */
    memcpy(out_session_data->ms_enc, master_key, sizeof(master_key));

    ret = 0;

finish:
    return ret;
}


/**
 * @brief  generate user DH key pair and
 *          encrypt the Master Key using PDH public key
 *
 * @param  [out]  out_session_data     The output session data
 * @param  [in]   pdh_cert             The hardware pdh cert
 *
 * @returns 0 for success, -1 for failure
 */
int  encrypt_master_key(session_buf* out_session_data, csv_cert* pdh_cert)
{
    int           ret            = -1;
    nonce_256     master_key     = {0};

    /* the exported hardware pdh public */
    ecc_point   hw_pdh_pubkey_coord = {0};
    EC_KEY     *hw_pdh_pubkey      = NULL;

    SM2CiphertextValue *cv = NULL;

    unsigned char      ms_enc[256]   = {0};
    size_t             ms_enc_len    = 256;

    unsigned char      *p  = NULL;
    const unsigned char      *buf  = (unsigned char*)ms_enc;

    if (get_public_key_in_cert(pdh_cert, &hw_pdh_pubkey_coord)) {
        printf("fail to read hardware pdh public key\n");
        goto finish;
    }

    /* it is plain master key in this moment */
    memcpy(master_key, out_session_data->ms_enc, sizeof(master_key));
    /* encrypt the master secret with the pdh public key */
    if (sm2_pubkey_coords_to_ec_key(&hw_pdh_pubkey_coord, &hw_pdh_pubkey) != 0) {
        printf("fail to coordinate to ec key\n");
        goto finish;
    }

    if (SM2_encrypt(NID_sm3, master_key, sizeof(master_key),
                    ms_enc, &ms_enc_len, hw_pdh_pubkey) == 0) {
        printf("fail to encrypt the master key\n");
        goto finish;
    }

    /* The output ms_enc is DER format,
     * manually convert it to C1||C2||C3 format
     */
    cv = d2i_SM2CiphertextValue(NULL, &buf, (long)ms_enc_len);
    if (cv == NULL) {
        printf("fail to convert encrypted ms to SM2CiphertextValue\n");
        goto finish;
    }

    /* copy SM2CiphertextValue to session data */
    memset(out_session_data->ms_enc, 0, sizeof(out_session_data->ms_enc));
    p = out_session_data->ms_enc;
    /* not compressed */
    *p = 0x4;
    p ++;

    /* copy random value coordinate */
    p += ECC_KEY_SIZE - BN_num_bytes(cv->xCoordinate);
    BN_bn2bin(cv->xCoordinate, p);
    p += BN_num_bytes(cv->xCoordinate);
    p += ECC_KEY_SIZE - BN_num_bytes(cv->yCoordinate);
    BN_bn2bin(cv->yCoordinate, p);
    p += BN_num_bytes(cv->yCoordinate);

    /* copy hash */
    memcpy(p, cv->hash->data, cv->hash->length);
    p += cv->hash->length;

    /* copy cipher */
    memcpy(p, cv->ciphertext->data, cv->ciphertext->length);

    ret = 0;

finish:
    SM2CiphertextValue_free(cv);
    return ret;
}


/*
 *
 *  calculate the virtual machine's digest
 *
 *  sign the session
 *
 *
 */
int  cal_session_mac(session_buf* out_session_data,
                     tek_tik    tek_tik_value,
                     guest_policy policy, uint32_t build_id)
{
    uint8_t session_mac_data[SESSION_MAC_SIZE];

    memset(session_mac_data, 0, sizeof(session_mac_data));

    if (build_id <= 1506) {
        sm3_hmac((unsigned char*)(&policy), sizeof(policy),
                 tek_tik_value.tik, sizeof(tek_tik_value.tik),
                 out_session_data->session_mac);
    } else {
        memcpy(session_mac_data, &policy, sizeof(policy));
        memcpy(session_mac_data + sizeof(policy),
               &out_session_data->ms_enc, SESSION_MAC_SIZE - sizeof(policy));
        sm3_hmac(session_mac_data, SESSION_MAC_SIZE,
                 tek_tik_value.tik, sizeof(tek_tik_value.tik),
                 (unsigned char*)(out_session_data->session_mac));
    }

    return 0;
}

int sign_session_buffer(session_buf* out_session_data, EC_KEY * guest_key_pair,
                        sm2_pubkey *guest_pubkey)
{
    int                ret                      = -1;
    unsigned char      sig[128]                 = {0};
    unsigned int       sig_len                  = 64;
    unsigned char      digest[EVP_MAX_MD_SIZE]  = {0};
    size_t             digest_len               = 32;
    ECDSA_SIG*         es                       = NULL;
    const uint8_t      *p                       = NULL;
    const BIGNUM       *br                      = NULL;
    const BIGNUM       *bs                      = NULL;
    unsigned char      buf[REAL_ECDSA_RS_SIZE];

    ret = SM2_compute_message_digest(EVP_sm3(), EVP_sm3(),
                   (unsigned char*)out_session_data,
                   sizeof(*out_session_data)-sizeof(out_session_data->sig),
                   (char*)guest_pubkey->userid.uid,
                   guest_pubkey->userid.len,
                   digest, &digest_len, guest_key_pair);

    if (ret == 0) {
        printf("SM2_compute_message_digest error\n");
        goto finish;
    }

    ret = SM2_sign(0, digest, (int)digest_len, sig, &sig_len, guest_key_pair);
    if (ret == 0) {
        printf("SM2_sign error\n");
        goto finish;
    }

    /*output signature*/
    p = sig;
    es = d2i_ECDSA_SIG(NULL, &p, sig_len);
    if(es == NULL) {
            printf("d2i_ECDSA_SIG error\n");
            goto finish;
    }

    ECDSA_SIG_get0(es, &br, &bs);
    if (BN_num_bytes(br) > REAL_ECDSA_RS_SIZE ||
        BN_num_bytes(bs) > REAL_ECDSA_RS_SIZE) {
        printf("invalid signature size\n");
        goto finish;
    }

    /* little endian in session_data */
    memset(buf, 0, sizeof(buf));
    BN_bn2bin(br, buf);
    reverse_bytes(buf, sizeof(buf));
    memcpy(out_session_data->sig, buf, REAL_ECDSA_RS_SIZE);

    memset(buf, 0, sizeof(buf));
    BN_bn2bin(bs, buf);
    reverse_bytes(buf, REAL_ECDSA_RS_SIZE);
    memcpy(out_session_data->sig + REAL_ECDSA_RS_SIZE, buf, REAL_ECDSA_RS_SIZE);

    ret = 0;

finish:
    return ret;
}

void  save_files(csv_cert *guest_owner_dh_cert, unsigned char *guest_priv_key,
                 tek_tik *tk, session_buf *out_session_data, nonce_256 *ms)
{
    char command[256]             = {0};

    if (!guest_owner_dh_cert || !guest_priv_key || !tk || !out_session_data) {
        printf("invalid data structure, fail to save to files\n");
        return ;
    }

    /* save the guest dh public key to cert file */
    write_file(GUEST_OWNER_DH_FILENAME, (unsigned char*)guest_owner_dh_cert, sizeof(csv_cert));
    printf("save guest owner dh cert raw data to %s\n", GUEST_OWNER_DH_FILENAME);
    /* save cert file to base64 format */
    sprintf(command, "base64 %s > %s", GUEST_OWNER_DH_FILENAME, GUEST_OWNER_DH_FILENAME_BASE64);
    system(command);
    printf("save guest owner dh cert base64 data to %s\n", GUEST_OWNER_DH_FILENAME_BASE64);

    /* save the guest dh private key to file */
    write_file(GUEST_PRIV_KEY_FILENAME, guest_priv_key, ECC_KEY_SIZE);
    printf("save guest owner dh private key data to %s\n", GUEST_PRIV_KEY_FILENAME);

    /* save the tektik key to file */
    write_file(GUEST_TEK_TIK_FILENAME, (unsigned char*)tk, sizeof(tek_tik));
    printf("save guest tek-tek key data to %s\n", GUEST_TEK_TIK_FILENAME);

    /* save session blob to file */
    write_file(LAUNCH_BLOB_FILENAME, (unsigned char*)out_session_data, sizeof(session_buf));
    printf("save guest launch blob raw data to %s\n", LAUNCH_BLOB_FILENAME);
    /* save session blob to base64 */
    sprintf(command, "base64 %s > %s", LAUNCH_BLOB_FILENAME, LAUNCH_BLOB_FILENAME_BASE64);
    system(command);
    printf("save guest launch blob base64 data to %s\n", LAUNCH_BLOB_FILENAME_BASE64);

    /* save master key to file */
    write_file(GUEST_MASTER_KEY_FILENAME, (unsigned char*)ms, sizeof(*ms));
    printf("save guest master key to %s\n", GUEST_MASTER_KEY_FILENAME);
}


int main(int argc, char *argv[])
{
    int opt;
    int option_index              = 0;
    unsigned long build_id        = ULONG_MAX;
    char *endptr                  = NULL;
    char *bios                    = NULL;
    char *kernel                  = NULL;
    char *cmdline                 = NULL;
    char *initrd                  = NULL;
    char *vcpu_model_arg          = NULL;
    unsigned long vcpu_num        = 0;
    uint8_t vm_digest[HASH_SIZE]  = {0};
    csv_cert   pdh_cert           = {0};
    session_buf out_session_data  = {0};
    char  * vm_id                 = "VM ID";
    char  * vm_version            = "VM 1.0";
    char  * user_id               = "user DH key";
    char  * user_data             = "user data";
    guest_policy    policy ;
    char  * policy_str            = NULL;
    nonce_256      master_key     = {0};
    /* guest's tek tik */
    /* guest's tek tik */
    tek_tik    tek_tik_value;
    char  *guest_priv_key_file    = NULL;

    /* guest's DH key */
    EC_KEY*            guest_key_pair     = NULL;
    unsigned char      guest_priv_key[ECC_KEY_SIZE] = {0};
    sm2_pubkey         guest_pubkey;

    /* csv cert structue storing the guest public key */
    csv_cert   guest_owner_dh_cert = {0};

    uint64_t  vcpu_model_value = 0;

    int ret                       = -1;

    if (argc <= 1) {
        usage();
        return -1;
    }

    while ((opt = getopt_long(argc, argv, "", long_args_options, &option_index)) != -1) {
        switch (opt) {
            case OPT_HELP:
                usage();
                return 0;
            case OPT_VERBOSE:
            case 'v':
                printf("Debug log print enabled.\n");
                break;
            case OPT_BUILD_ID:
                build_id = strtoul(optarg, &endptr, 10);
                if (*endptr != '\0') {
                    printf("Wrong FW build ID: %s\n", optarg);
                    return -1;
                }
                printf("build id %lu\n", build_id);
                break;
            case OPT_BIOS:
                if (optarg == NULL || optarg[0] == '-') {
                    usage();
                    return -1;
                }
                bios = optarg;
                printf("bios file = %s\n", bios);
                break;
            case OPT_KERNEL:
                kernel = optarg;
                printf("kernel file = %s\n", kernel);
                break;
            case OPT_CMDLINE:
                cmdline = optarg;
                printf("cmdline = '%s'\n", cmdline);
                break;
            case OPT_INITRD:
                initrd = optarg;
                printf("initrd = %s \n", initrd);
                break;
            case OPT_VM_ID:
                vm_id =  optarg;
                printf("vm id = %s \n", vm_id);
                break;
            case OPT_VM_VERSION:
                vm_version = optarg;
                printf("vm version = %s \n", vm_version);
                break;
            case OPT_VCPU_MODEL:
                vcpu_model_arg = optarg;
                if (good_vcpu_model(vcpu_model_arg) == -1)
                    return -1;
                printf("vcpu_model = %s \n", vcpu_model_arg);
               break;
            case OPT_VCPU_NUM:
                vcpu_num = strtoul(optarg, &endptr, 10);
                if (*endptr != '\0') {
                    printf("Wrong vcpu num: %s\n", optarg);
                    return -1;
                }
                printf("vcpu num = %lu \n", vcpu_num);
                break;
            case OPT_USER_ID:
                user_id = optarg;
                printf("user id = %s \n", user_id);
                break;
            case OPT_USER_DATA:
                user_data = optarg;
                printf("user data = %s \n", user_data);
                break;
            case OPT_GUEST_POLICY:
                policy_str = optarg;
                break;
            case OPT_GUEST_PRIV_KEY:
                guest_priv_key_file = optarg;
                printf("guest private key file = %s \n", guest_priv_key_file);
                break;
            default:
                usage();
                return -1;
        }
    }

    if (policy_str == NULL) {
        printf("no policy input\n");
        usage();
        return -1;
    } else {
        long int v;
        if (strlen(policy_str) >= 2 && policy_str[0] == '0' &&
            (policy_str[1] == 'x' || policy_str[1] == 'X')) {
            v = strtol(policy_str, NULL, 16);
        } else {
            v = strtol(policy_str, NULL, 10);
        }
        policy = *(guest_policy*)&v;
        printf("guest policy 0x%x\n", *(int*)&policy);
    }

    if (build_id == ULONG_MAX) {
        printf(" Assume build id is 1545 or larger\n");
        build_id = 1545;
    }

    if (guest_priv_key_file == NULL) {
        printf("no guest DH private key file input\n");
        usage();
        return -1;
    }

    memset(&guest_owner_dh_cert, 0, sizeof(guest_owner_dh_cert));
    /*
     * create the user's dh key pair
     */
    memset(&guest_pubkey, 0, sizeof(guest_pubkey));
    guest_key_pair = get_sm2_key_pair(guest_priv_key_file, &guest_pubkey, guest_priv_key);
    if (guest_key_pair == NULL) {
        printf("fail to generate guest owner dh key pair\n");
        goto finish;
    }

    memcpy(guest_pubkey.userid.uid, user_id, strlen(user_id));
    guest_pubkey.userid.len = strlen(user_id);

    strncpy((char*)out_session_data.vm_id, vm_id, sizeof(out_session_data.vm_id));
    strncpy((char*)out_session_data.vm_version, vm_version,
            sizeof(out_session_data.vm_version));
    strncpy((char*)out_session_data.user_data, user_data,
            sizeof(out_session_data.user_data));

    vcpu_model_value = get_vcpu_model_value(vcpu_model_arg);
    if (vcpu_model_value == 0) {
        printf("vcpu model value is 0\n");
        goto finish;
    }

    if (calculate_vm_digest( policy, bios, kernel, cmdline, initrd,
                            vcpu_model_value, vcpu_num,
                            vm_digest) != 0) {
        printf("fail to calculate VM digest\n");
        goto finish;
    }

    /* copy vm disgest to session data */
    memcpy(out_session_data.vm_digest, vm_digest, sizeof(vm_digest));

    /* create tek/tik and save to session data */
    if (generate_tk(&out_session_data, &tek_tik_value, &master_key, build_id)) {
        printf("create and fill tek/tik error\n");
        goto finish;
    }

    if (read_file(PDH_FILENAME, (unsigned char*)&pdh_cert, sizeof(csv_cert)) != 0) {
        printf("read pdh cert file error\n");
        goto finish;
    }

    if (encrypt_master_key(&out_session_data, &pdh_cert)) {
        printf("fail to encrypt master key\n");
        goto finish;
    }

    if (cal_session_mac(&out_session_data, tek_tik_value,policy, build_id)) {
        printf("fail to calculate the session mac\n");
        goto finish;
    }

    if (sign_session_buffer(&out_session_data, guest_key_pair, &guest_pubkey)) {
        printf("fail to sign session buffer\n");
        goto finish;
    }

    /*
     * session data structure is all ready.
     *
     * save tek_tik, guest pubkey, guest privkey, blob to file
     */
    if (fill_pubkey_in_cert(&guest_owner_dh_cert, &guest_pubkey) != 0){
        printf("fail to write guest owner dh public key to cert\n");
        goto finish;
    }
    guest_owner_dh_cert.version = 1;// default version

    save_files(&guest_owner_dh_cert, guest_priv_key, &tek_tik_value,
               &out_session_data, &master_key);

    ret = 0;

finish:

    return ret;
}
